#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
分析pkl文件结构，查看是否遗漏了轨迹相关的数据
"""

import os
import pickle
import numpy as np
import argparse

def analyze_pkl_file(pkl_path):
    """详细分析单个pkl文件的结构"""
    print(f"\n{'='*60}")
    print(f"分析文件: {os.path.basename(pkl_path)}")
    print(f"{'='*60}")
    
    try:
        with open(pkl_path, 'rb') as f:
            data = pickle.load(f)
        
        print(f"文件大小: {os.path.getsize(pkl_path)} bytes")
        print(f"顶级键值: {list(data.keys())}")
        
        for key, value in data.items():
            if key == 'pred_output_list':
                print(f"\n{key}: 长度={len(value)}")
                if len(value) > 0 and value[0] is not None:
                    pred_output = value[0]
                    print(f"  第一个预测输出的键值:")
                    for pred_key, pred_value in pred_output.items():
                        if isinstance(pred_value, np.ndarray):
                            print(f"    {pred_key}: {pred_value.shape} ({pred_value.dtype})")
                        elif pred_value is None:
                            print(f"    {pred_key}: None")
                        else:
                            print(f"    {pred_key}: {type(pred_value).__name__} = {pred_value}")
            else:
                if isinstance(value, np.ndarray):
                    print(f"{key}: {value.shape} ({value.dtype})")
                elif isinstance(value, (list, tuple)):
                    print(f"{key}: {type(value).__name__}(长度={len(value)})")
                else:
                    print(f"{key}: {type(value).__name__} = {value}")
        
        # 查找可能的轨迹相关字段
        trajectory_related_keys = []
        if 'pred_output_list' in data and len(data['pred_output_list']) > 0:
            pred_output = data['pred_output_list'][0]
            if pred_output:
                for key in pred_output.keys():
                    if any(word in key.lower() for word in ['traj', 'path', 'motion', 'transl', 'joints', 'vertices']):
                        trajectory_related_keys.append(key)
        
        if trajectory_related_keys:
            print(f"\n轨迹相关字段:")
            for key in trajectory_related_keys:
                if pred_output[key] is not None:
                    if isinstance(pred_output[key], np.ndarray):
                        print(f"  {key}: {pred_output[key].shape} ({pred_output[key].dtype})")
                    else:
                        print(f"  {key}: {type(pred_output[key]).__name__}")
                else:
                    print(f"  {key}: None")
        else:
            print("\n未找到明显的轨迹相关字段")
            
        return True
        
    except Exception as e:
        print(f"错误: {e}")
        return False

def compute_trajectory_from_data(pred_output):
    """从pred_output数据计算可能的轨迹信息"""
    trajectories = {}
    
    # 1. 根部位置轨迹（从平移数据）
    if 'pred_transl' in pred_output and pred_output['pred_transl'] is not None:
        trajectories['root_trajectory'] = pred_output['pred_transl']
    
    # 2. 全局朝向轨迹
    if 'pred_global_orient' in pred_output and pred_output['pred_global_orient'] is not None:
        trajectories['global_orient_trajectory'] = pred_output['pred_global_orient']
    
    # 3. 关节轨迹
    joint_keys = ['pred_body_joints_3d', 'pred_lhand_joints_3d', 'pred_rhand_joints_3d']
    for key in joint_keys:
        if key in pred_output and pred_output[key] is not None:
            trajectories[f"{key}_trajectory"] = pred_output[key]
    
    # 4. 顶点轨迹（可能太大，但也是轨迹的一种）
    if 'pred_vertices_smpl' in pred_output and pred_output['pred_vertices_smpl'] is not None:
        trajectories['vertices_trajectory'] = pred_output['pred_vertices_smpl']
    
    # 5. 手部位置轨迹
    hand_orient_keys = [
        'left_hand_global_orient_body', 'right_hand_global_orient_body',
        'left_hand_global_orient_hand', 'right_hand_global_orient_hand'
    ]
    for key in hand_orient_keys:
        if key in pred_output and pred_output[key] is not None:
            trajectories[f"{key}_trajectory"] = pred_output[key]
    
    return trajectories

def analyze_trajectory_potential(pkl_dir):
    """分析整个目录的轨迹潜力"""
    pkl_files = sorted([f for f in os.listdir(pkl_dir) if f.endswith('.pkl')])
    
    if not pkl_files:
        print(f"在 {pkl_dir} 中没有找到pkl文件")
        return
    
    print(f"找到 {len(pkl_files)} 个pkl文件")
    
    # 分析前几个文件
    for i, pkl_file in enumerate(pkl_files[:3]):
        pkl_path = os.path.join(pkl_dir, pkl_file)
        analyze_pkl_file(pkl_path)
        
        if i == 0:  # 只对第一个文件计算轨迹
            try:
                with open(pkl_path, 'rb') as f:
                    data = pickle.load(f)
                
                if 'pred_output_list' in data and len(data['pred_output_list']) > 0:
                    pred_output = data['pred_output_list'][0]
                    if pred_output:
                        trajectories = compute_trajectory_from_data(pred_output)
                        print(f"\n从数据中可计算的轨迹:")
                        for traj_key, traj_data in trajectories.items():
                            if isinstance(traj_data, np.ndarray):
                                print(f"  {traj_key}: {traj_data.shape} ({traj_data.dtype})")
            except Exception as e:
                print(f"计算轨迹时出错: {e}")
    
    print(f"\n{'='*60}")
    print("总结:")
    print("1. 如果要包含轨迹数据，主要字段包括:")
    print("   - pred_transl: 根部位置轨迹")
    print("   - pred_global_orient: 全局朝向轨迹")
    print("   - pred_body_joints_3d: 身体关节轨迹")
    print("   - pred_lhand_joints_3d: 左手关节轨迹")
    print("   - pred_rhand_joints_3d: 右手关节轨迹")
    print("   - pred_vertices_smpl: 顶点轨迹（较大）")
    print("2. 这些数据在合并时形成时间序列，即构成动画轨迹")
    print(f"{'='*60}")

def main():
    parser = argparse.ArgumentParser(description='分析pkl文件结构，查找轨迹数据')
    parser.add_argument('--input', '-i',
                       default='./mocap_output_full/mocap',
                       help='包含pkl文件的目录')
    parser.add_argument('--file', '-f',
                       help='分析单个pkl文件（可选）')
    
    args = parser.parse_args()
    
    if args.file:
        if os.path.exists(args.file):
            analyze_pkl_file(args.file)
        else:
            print(f"文件不存在: {args.file}")
    else:
        if os.path.exists(args.input):
            analyze_trajectory_potential(args.input)
        else:
            print(f"目录不存在: {args.input}")

if __name__ == "__main__":
    main()